# -*- coding: utf-8 -*-
# 导入sklearn包的决策树类
from sklearn import tree

#准备训练数据
features=[[210,2],[120,240],[150,300],[130,200],[110,150],[140,100],[160,50],[170,30],[180,220],[190,10]]
labels=['动作片','爱情片','爱情片','爱情片','爱情片','动作片','动作片','动作片','爱情片','动作片']

# 创建决策树分类器
clf = tree.DecisionTreeClassifier()
#训练数据并得到模型
model= clf.fit(features, labels)
#使用模型预测新数据

result = clf.predict([[245,242]])
print(result)  # 输出预测结果
# 导出dot文件
with open("firstmldemo.dot", 'w') as f:
    f = tree.export_graphviz(clf, out_file=f, feature_names=['时长', '评分'], class_names=labels, filled=True, rounded=True, special_characters=True)   